from RAG import rag
from utils import Agent_role
import os
import re


def framework_processing(text):
    pattern = r'第[一二三四五六七八九十\d]+章'
    text_list = [i for i in text.split('\n') if i]
    index_list = [index for index, i in enumerate(text_list) if re.findall(pattern, i)] + [len(text_list)]
    new_text_list = [text_list[index_list[i]:index_list[i + 1]] for i in list(range(len(index_list) - 1))]
    for index, i in enumerate(new_text_list):
        new_text_list[index] = {'title': i[0], 'title_content': ''.join(i[1:])}
    return new_text_list


if __name__ == '__main__':
    question = '写一篇游戏相关题材的长篇小说'
    agent = Agent_role()
    new_question = agent.problem_norm(question)
    print("问题规范:", new_question)
    file_path = './database/data.pkl'  # 外部知识向量，缺乏向量数据库知识，暂时用pkl文件代替数据库，需要时加载该文件
    if os.path.exists(file_path):
        # 无知识库
        knowledge_list = rag(new_question, file_path)
        #利用语料构建提示词，暂待更新
    flow = agent.problem_architecture(new_question)
    print("小说基本框架:", flow)
    novel_arch = framework_processing(flow)
    word = 10000
    word_len = (int(word / len(novel_arch) / 1000) + 1) * 1000
    novel_text = ''
    novel_above = '小说开头：'
    for index, i in enumerate(novel_arch):
        if i['title_content']:
            if index == len(novel_arch) - 1:
                novel_above += '小说结局：'
            novel = agent.expanding(novel_above, i['title_content'], word_len)
            novel_text += i['title'] + '\n' + re.sub('\n\n', '\n', novel) + '\n'
            novel_above += i['title_content']
    novel_above = re.sub('小说开头：|小说结尾：', '', novel_above)
    novel_path = './novel.txt'
    with open(novel_path, 'w', encoding = 'utf-8') as file:
        file.write(novel_text)
